# nn.MaxPool2d()
# 函数通过对图像进行最大值下采样，可以将图像尺寸便小。与之相对应得nn.MaxUnpool2d()
# 函数可以实现对图像进行上采样，将图像尺寸变大。
#
#
#
# 请查询nn.MaxUnpool2d()
# 函数的使用手册，并设计一个先降维后升维的卷积神经网络(u - net)。
#
#
#
# 1.
# 该网络执行如下 ** 降维 ** 操作：
#
# - 第一层： 使用8个大小为3×3、步长为1、填充为1的卷积核，将输入1×224×224
# 的图像输出为8×224×224
# 的图像， 使用ReLU为激活函数;
#
# - 第二层： 大小为2×2、 步长为2、无填充的极大值池化层，将8×224×224
# 的图像输出为8×112×112;
#
# - 第三层： 使用16个大小为3×3、步长为1、填充为1的卷积核，将输入8×112×112
# 的图像输出为16×112×112
# 的图像， 使用ReLU为激活函数;
#
# - 第四层： 大小为2×2、步长为2、无填充的极大值池化层，将16×112×112
# 的图像输出为16×56×56;
#
# - 第五层： 使用32个大小为3×3、步长为1、填充为1的卷积核，将输入16×56×56
# 的图像输出为32×56×56
# 的图像， 使用ReLU为激活函数;
#
# - 第六层： 大小为2×2、步长为2、填充为1的极大值池化层，将32×56×56
# 的图像输出为32×28×28;
#
# 2.
# 该网络接着执行如下 ** 升维 ** 操作：
#
# - 第七层： 大小为2×2、步长为2、填充为1的 ** 反极大值池化层 **，将32×28×28
# 的图像输出为32×56×56;
#
# - 第八层： 使用16个大小为3×3、步长为1、填充为1的卷积核，将输入32×56×56
# 的图像输出为16×56×56
# 的图像， 使用ReLU为激活函数;
#
# - 第九层： 大小为2×2、步长为2、填充为1的 ** 反极大值池化层 **，将16×56×56
# 的图像输出为16×112×112;
#
# - 第十层： 使用8个大小为3×3、步长为1、填充为1的卷积核，将输入16×112×112
# 的图像输出为8×112×112
# 的图像， 使用ReLU为激活函数;
#
# - 第十一层： 大小为2×2、步长为2、填充为1的 ** 反极大值池化层 **，将8×112×112
# 的图像输出为8×224×224;
#
# - 第十二层： 使用1个大小为3×3、步长为1、填充为1的卷积核，将输入8×224×224
# 的图像输出为1×224×224
# 的图像， 使用ReLU为激活函数;



import torch
import torch.nn as nn

# 定义down_net按要求实现图像卷积和下采样
down_net = nn.Sequential(
    nn.Conv2d(1, 8, kernel_size=3, stride=1, padding=1),
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=2, stride=2),
    nn.Conv2d(8, 16, kernel_size=3, stride=1, padding=1),
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=2, stride=2),
    nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1),
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=2, stride=2)
)

# 定义up_net按要求实现图像上采样和卷积
up_net = nn.Sequential(
    nn.ConvTranspose2d(32, 32, kernel_size=2, stride=2),
    nn.Conv2d(32, 16, kernel_size=3, stride=1, padding=1),
    nn.ReLU(),
    nn.ConvTranspose2d(16, 16, kernel_size=2, stride=2),
    nn.Conv2d(16, 8, kernel_size=3, stride=1, padding=1),
    nn.ReLU(),
    nn.ConvTranspose2d(8, 8, kernel_size=2, stride=2),
    nn.Conv2d(8, 1, kernel_size=3, stride=1, padding=1),
    nn.ReLU()
)

# 测试网络down_net是否符合按要求将图像进行下采样：
X = torch.randn(1, 1, 224, 224)  # 测试输入数据
print('降维过程：')
for layer in down_net:
    X = layer(X)
    print(layer.__class__.__name__, 'output shape: \t', X.shape)

# 测试网络up_net是否按要求将图像进行上采样：
print('升维过程：')
for layer in up_net:
    X = layer(X)
    print(layer.__class__.__name__, 'output shape: \t', X.shape)



# # 代码块的输出需与下面结果保持一致：
# """
# 降维过程：
# Conv2d output shape: 	 torch.Size([1, 8, 224, 224])
# ReLU output shape: 	 torch.Size([1, 8, 224, 224])
# MaxPool2d_Indices output shape: 	 torch.Size([1, 8, 112, 112])
# Conv2d output shape: 	 torch.Size([1, 16, 112, 112])
# ReLU output shape: 	 torch.Size([1, 16, 112, 112])
# MaxPool2d_Indices output shape: 	 torch.Size([1, 16, 56, 56])
# Conv2d output shape: 	 torch.Size([1, 32, 56, 56])
# ReLU output shape: 	 torch.Size([1, 32, 56, 56])
# MaxPool2d_Indices output shape: 	 torch.Size([1, 32, 28, 28])
# 升维过程：
# MaxUnpool2d_Indices output shape: 	 torch.Size([1, 32, 56, 56])
# Conv2d output shape: 	 torch.Size([1, 16, 56, 56])
# ReLU output shape: 	 torch.Size([1, 16, 56, 56])
# MaxUnpool2d_Indices output shape: 	 torch.Size([1, 16, 112, 112])
# Conv2d output shape: 	 torch.Size([1, 8, 112, 112])
# ReLU output shape: 	 torch.Size([1, 8, 112, 112])
# MaxUnpool2d_Indices output shape: 	 torch.Size([1, 8, 224, 224])
# Conv2d output shape: 	 torch.Size([1, 1, 224, 224])
# ReLU output shape: 	 torch.Size([1, 1, 224, 224])
# """